# -*- coding: utf-8 -*-
# @Time : 2021-11-17 18:49
# @Author : lwb
# @Site : 
# @File : vocab.py
from collections import defaultdict
# 按字建立字典（不是词语）
class Vocab:
    def __init__(self,tokens=None):
        self.idx_to_token=list()
        #使用列表存储所有的标记，从而根据索引值获取相应 的标记
        self.token_to_idx=dict()
        if tokens is not None:
            if "<unk>" not in tokens:
                tokens = tokens + ["<unk>"]
            for token in tokens:
                self.idx_to_token.append(token)
                self.token_to_idx[token] = len(self.idx_to_token) - 1
            self.unk = self.token_to_idx['<unk>']

    @classmethod
    def build(cls,text,min_freq=1,reserved_tokens=None):
        # 创建此表，输入的text包含若干句子，每个句子由若干标记构成
        token_freqs=defaultdict(int)  # 存储标记及出现次数的映射字典
        for sentence in text:
            for token in sentence:
                token_freqs[token]+=1
        # 未登录词 unknoen word 标记<unk>
        uniq_tokens = ["<unk>"] + (reserved_tokens if reserved_tokens else [])
        uniq_tokens += [token for token, freq in token_freqs.items() if freq >= min_freq and token != "<unk>"]
        return cls(uniq_tokens)

    def __len__(self):
        # 返回此表的大小，即此表中有多少个互不相同的标记
        return len(self.idx_to_token)
    def __getitem__(self, token):
        # 查找输入标记对应的索引值，如果该值不存在，则返回标记<unk> 的索引值0
        return self.token_to_idx.get(token,self.unk)
    def convert_tokens_to_ids(self,tokens):
        # 查找一系列输入标记对应的索引值
        return [self[token] for token in tokens]
    def convert_ids_to_tokens(self,indices):
        # 查找一系列索引值对应的标记
        return [self.idx_to_token[index] for index in indices]
